package meiyunting.shiwuzhaoling.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 自定义负载均衡策略
 * 实现加权轮询，可以根据实例的权重分配请求
 */
@Slf4j
public class CustomLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    // 计数器，记录当前轮询位置
    private final AtomicInteger position;
    
    // 服务名称
    private final String serviceId;
    
    // 服务实例列表提供者
    private final ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;
    
    // 权重缓存
    private final Map<String, Integer> weightCache = new ConcurrentHashMap<>();
    
    // 随机数生成器，用于在相同权重的情况下随机选择
    private final Random random = new Random();
    
    public CustomLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider, 
                              String serviceId) {
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.serviceId = serviceId;
        this.position = new AtomicInteger(0);
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider.getIfAvailable();
        return supplier.get(request).next()
                .map(serviceInstances -> processInstanceResponse(serviceInstances, request));
    }

    private Response<ServiceInstance> processInstanceResponse(List<ServiceInstance> instances, Request request) {
        if (instances.isEmpty()) {
            log.warn("【负载均衡】没有可用的服务实例: {}", serviceId);
            System.out.println("【控制台输出】负载均衡警告: 没有可用的服务实例: " + serviceId);
            return new EmptyResponse();
        }
        
        log.info("【负载均衡】开始为服务 {} 选择实例，当前可用实例数: {}", serviceId, instances.size());
        System.out.println("【控制台输出】开始为服务 " + serviceId + " 选择实例，当前可用实例数: " + instances.size());
        
        // 根据权重选择实例
        ServiceInstance selectedInstance = selectInstanceByWeight(instances);
        
        log.info("【负载均衡】已选择实例: {}:{}, 实例ID: {}", 
                selectedInstance.getHost(), selectedInstance.getPort(), selectedInstance.getInstanceId());
        System.out.println("【控制台输出】已选择实例: " + selectedInstance.getHost() + ":" + 
                selectedInstance.getPort() + ", 实例ID: " + selectedInstance.getInstanceId());
        
        return new DefaultResponse(selectedInstance);
    }
    
    private ServiceInstance selectInstanceByWeight(List<ServiceInstance> instances) {
        // 如果只有一个实例，直接返回
        if (instances.size() == 1) {
            ServiceInstance instance = instances.get(0);
            log.info("【负载均衡】只有一个可用实例: {}:{}", instance.getHost(), instance.getPort());
            System.out.println("【控制台输出】只有一个可用实例: " + instance.getHost() + ":" + instance.getPort());
            return instance;
        }
        
        // 获取总权重
        int totalWeight = 0;
        StringBuilder weightInfo = new StringBuilder("【实例权重信息】");
        for (ServiceInstance instance : instances) {
            int weight = getInstanceWeight(instance);
            totalWeight += weight;
            weightInfo.append(String.format(" %s:%s(权重:%d)", 
                    instance.getHost(), instance.getPort(), weight));
        }
        log.info(weightInfo.toString());
        System.out.println("【控制台输出】" + weightInfo.toString());
        
        // 如果所有实例权重都是1，使用简单轮询
        if (totalWeight == instances.size()) {
            int pos = Math.abs(this.position.incrementAndGet()) % instances.size();
            ServiceInstance instance = instances.get(pos);
            log.info("【负载均衡】使用轮询策略，选择第 {} 个实例: {}:{}", 
                    pos + 1, instance.getHost(), instance.getPort());
            System.out.println("【控制台输出】使用轮询策略，选择第 " + (pos + 1) + 
                    " 个实例: " + instance.getHost() + ":" + instance.getPort());
            return instance;
        }
        
        // 加权选择
        int randomWeight = random.nextInt(totalWeight);
        int current = 0;
        for (ServiceInstance instance : instances) {
            int weight = getInstanceWeight(instance);
            current += weight;
            if (randomWeight < current) {
                log.info("【负载均衡】使用加权策略，选择实例: {}:{} (权重:{}/总权重:{})", 
                        instance.getHost(), instance.getPort(), weight, totalWeight);
                System.out.println("【控制台输出】使用加权策略，选择实例: " + instance.getHost() + ":" + 
                        instance.getPort() + " (权重:" + weight + "/总权重:" + totalWeight + ")");
                return instance;
            }
        }
        
        // 如果出现问题，随机选择一个
        ServiceInstance instance = instances.get(random.nextInt(instances.size()));
        log.warn("【负载均衡】加权选择出现问题，随机选择实例: {}:{}", instance.getHost(), instance.getPort());
        System.out.println("【控制台输出】加权选择出现问题，随机选择实例: " + 
                instance.getHost() + ":" + instance.getPort());
        return instance;
    }
    
    /**
     * 获取实例权重
     * 默认权重为1，可以通过metadata中的weight属性设置权重
     */
    private int getInstanceWeight(ServiceInstance instance) {
        String instanceId = instance.getInstanceId();
        if (weightCache.containsKey(instanceId)) {
            return weightCache.get(instanceId);
        }
        
        // 从metadata中获取权重，默认为1
        int weight = 1;
        if (instance.getMetadata().containsKey("weight")) {
            try {
                weight = Integer.parseInt(instance.getMetadata().get("weight"));
                if (weight <= 0) {
                    weight = 1;
                }
            } catch (NumberFormatException e) {
                // 解析失败，使用默认权重
                log.warn("【负载均衡】实例 {} 的权重解析失败，使用默认权重 1", instanceId);
                System.out.println("【控制台输出】实例 " + instanceId + " 的权重解析失败，使用默认权重 1");
            }
        }
        
        weightCache.put(instanceId, weight);
        return weight;
    }
} 